import asyncio
import json
import os
import random

from typing import List

from agents import Agent, RunContextWrapper, Runner, function_tool
from agents.models.openai_provider import OpenAIProvider

from game_helpers import get_random_card_list, get_cached_expressions


API_KEY = os.getenv("EXAMPLE_API_KEY")
BASE_URL = os.getenv("EXAMPLE_BASE_URL")

provider = OpenAIProvider(
        api_key=API_KEY,
        base_url=BASE_URL,
        use_responses=False
    )
model = provider.get_model("qwen2.5")


instruction_template_dict = {
    "GameDealer": """
        Generate 4 random natural numbers between 1 and 13, include 1 and 13. Just return 4 numbers in an array, don't include other content. The returned array should not be repeated with the following arrays:
        {old_arrays}
        """,
    "MathProdigy": """You are a helpful agent. """,
    "GameJudger": """You are a helpful agent. """,
    "GamePlayer": """You are a helpful agent. """,
}


user_prompt_template_dict = {
    "GameDealer": """generate an array""",
    "MathProdigy": """
        What's the 24 points expression of {last_cards_posted} ? 
        If the result is 'expression not found', just return 'expression not found'. 
        If the result is an arithmetic expression, just return the expression itself and do not add anything else.
        """,
    "GameJudger": """
        Cards posted is '{last_cards_posted}', what's the check result of {expression} ? Just return the check result itself such as 'Correct' or 'Wrong', and do not add anything else such as 'The check result is ...'. 
    """,
    "GamePlayer": """
        What's the human user's reply of {last_cards_posted} ? Just return the reply itself as is, please follow below format:
        "reply": "..."
    """,
}


class CustomContext:
    def __init__(self, context_variables: dict):
        self.context_variables = context_variables
        

def get_instruction(
    run_context: RunContextWrapper[CustomContext], agent: Agent[CustomContext]
) -> str:
    global instruction_template_dict

    context = run_context.context
    context_variables = context.context_variables

    agent_name = context_variables["agent_name"]
    instruction_template = instruction_template_dict[agent_name]
    instruction = instruction_template

    if agent_name == "GameDealer":
        last_cards_posted = context_variables["old_arrays"]
        instruction = instruction_template.format(old_arrays=last_cards_posted)
   
    return instruction


def get_user_prompt(context_variables):
    global user_prompt_template_dict

    agent_name = context_variables["agent_name"]
    user_prompt_template = user_prompt_template_dict[agent_name]
    user_prompt = user_prompt_template

    if agent_name == "MathProdigy":
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(last_cards_posted=last_cards_posted)

    elif agent_name == "GameJudger":
        expression = context_variables["expression"]
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(expression=expression,last_cards_posted=last_cards_posted)

    elif agent_name == "GamePlayer":
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(last_cards_posted=last_cards_posted)

    return user_prompt


@function_tool
def get_24_points_expression_func(last_cards_posted: str) -> str:
    """Resolve the expression of 24 points game, return an arithmetic expression.

    Keyword arguments:
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    point_list = json.loads(last_cards_posted)  
    if len(point_list) == 0:
        return "expression not found"

    expressions = get_cached_expressions(point_list)
    
    result = "expression not found"
    if len(expressions) > 0:
        random_idx = random.randint(0, len(expressions)-1)
        expression = f"'{expressions[random_idx]}'".replace("'", "")
        print(f"The resolved 24 points expression is '{expression}'")
        return expression
    
    return "expression not found"


@function_tool
def check_24_points_expression_func(expression: str, last_cards_posted: str) -> str:
    """Check if the result of an arithmetic expression is equal 24, return 'Correct' or 'Wrong'.

    Keyword arguments:
      expression: an arithmetic expression
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    result = eval(expression.replace("'", ""))
    if abs(result - 24) < 0.001:
        return "Correct"

    return "Wrong"


@function_tool
def get_human_reply_func(last_cards_posted: str) -> str:
    """Interact with a human user and get a reply from the human user. The replay should be 'deal', 'help', 'exit' or an an arithmetic expression.

    Keyword arguments:
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    PROMPT_TEMPLATE: str = """
    Cards the dealer just posted: {content}
    Please give an expression for the four operations that results in 24.
    Type 'help' if you feel it's difficult.
    Type 'deal' if you want the dealer to deal cards again.
    Type 'exit' if you want to exit this game, type 'exit'.
    """

    point_list = json.loads(last_cards_posted)
    card_list = get_random_card_list(point_list)
    cards_content = f"{{'card_list': {card_list}, 'point_list': {point_list}}}"

    prompt = PROMPT_TEMPLATE.format(content=cards_content)
    human_reply = input(prompt)

    return {"human_reply": f"{human_reply}"}


agent_bill = Agent(
    name="GameDealer",
    instructions=get_instruction,
    model=model,
    tools=[],
)


agent_gauss = Agent(
    name="MathProdigy",
    instructions=get_instruction,
    model=model,
    tools=[get_24_points_expression_func],
)


agent_peter = Agent(
    name="GameJudger",
    instructions=get_instruction,
    model=model,
    tools=[check_24_points_expression_func],
)


agent_david = Agent(
    name="GamePlayer",
    instructions=get_instruction,
    model=model,
    tools=[get_human_reply_func],
)


async def deal_cards(old_arrays: List[int]) -> str:
    global agent_bill

    print(f"used old_arrays is :{old_arrays}")
    context_var_dict = {
        "agent_name":"GameDealer", 
        "old_arrays": f"{old_arrays}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_bill,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    cards_posted = response.final_output
    print(cards_posted)
    return cards_posted


async def machine_give_expression(last_cards_posted: str) -> str:
    global agent_gauss

    context_var_dict = {
        "agent_name":"MathProdigy", 
        "last_cards_posted": f"{last_cards_posted}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_gauss,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    expression = response.final_output
    print(expression)
    return expression


async def check_expression(expression: str, last_cards_posted: str) -> str:
    global agent_peter

    context_var_dict = {
        "agent_name":"GameJudger", 
        "expression": f"{expression}",
        "last_cards_posted": f"{last_cards_posted}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_peter,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    check_result = response.final_output
    print(check_result)
    return check_result


async def get_human_reply(last_cards_posted: str) -> str:
    global agent_david

    context_var_dict = {
        "agent_name":"GamePlayer", 
        "last_cards_posted": f"{last_cards_posted}"
    }
    
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_david,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    human_reply = response.final_output.replace('reply', '').replace('.', '').replace(':', '').replace('(', '').replace(')', '').replace('"', '').strip()
    print(human_reply)
    return human_reply


async def main_func():
    
    old_arrays = []
    last_cards_posted = await deal_cards(old_arrays)

    while True:
        human_reply = await get_human_reply(last_cards_posted)
        print(f"human_reply is '{human_reply}'")

        if human_reply == "deal":
            old_arrays.append(json.loads(last_cards_posted))
            last_cards_posted = await deal_cards(old_arrays)
            continue
        elif human_reply == "help":
            expression = await machine_give_expression(last_cards_posted)
        elif human_reply == "exit":
            print("Bye bye, have a good day!")
            break
        else:
            expression = human_reply
        
        if expression != "expression not found":
            check_result = await check_expression(expression, last_cards_posted)
        else:
            check_result = "Correct"
                        
        if check_result == "Correct":
            old_arrays.append(json.loads(last_cards_posted))
            last_cards_posted = await deal_cards(old_arrays)


if __name__ == "__main__":
    asyncio.run(main_func())

